package net.semanticmetadata.lire;

import com.hd.model.Pic;
import net.semanticmetadata.lire.builders.DocumentBuilder;
import net.semanticmetadata.lire.builders.GlobalDocumentBuilder;
import net.semanticmetadata.lire.imageanalysis.features.global.AutoColorCorrelogram;
import net.semanticmetadata.lire.imageanalysis.features.global.CEDD;
import net.semanticmetadata.lire.imageanalysis.features.global.FCTH;
import net.semanticmetadata.lire.searchers.GenericFastImageSearcher;
import net.semanticmetadata.lire.searchers.ImageSearchHits;
import net.semanticmetadata.lire.searchers.ImageSearcher;
import org.apache.lucene.analysis.core.WhitespaceAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.store.FSDirectory;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;

/**
 * Created by qingshengzheng on 17/5/16.
 */
public class HDLireOper {
    private GlobalDocumentBuilder globalDocumentBuilder;
    private IndexWriter iw;
    private IndexReader ir;
    private ImageSearcher searcher;
    private static final String indexpath="/work/lire/index";

    public HDLireOper() {
        initBuilder();
      //  initIndexWrite(indexpath);
        initImageSearch();
        initIndexReader(indexpath);
    }



    public void initImageSearch() {
        if (searcher != null) return;
        searcher = new GenericFastImageSearcher(30, CEDD.class);
    }

    public void initIndexReader(String index) {
        if (ir != null) return;
        try {
            ir = DirectoryReader.open(FSDirectory.open(Paths.get(index)));
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 初始化
     */
    private void initBuilder() {
        if (globalDocumentBuilder == null)
            this.globalDocumentBuilder = new GlobalDocumentBuilder(false, false);
        globalDocumentBuilder.addExtractor(CEDD.class);
        globalDocumentBuilder.addExtractor(FCTH.class);
        globalDocumentBuilder.addExtractor(AutoColorCorrelogram.class);
    }

    public  void initIndexWrite() {
     initIndexWrite(indexpath);
    }




    private void initIndexWrite(String index) {
        IndexWriterConfig conf = new IndexWriterConfig(new WhitespaceAnalyzer());

            try {
                iw = new IndexWriter(FSDirectory.open(Paths.get(index)), conf);
                iw.commit();
            } catch (IOException e) {
                e.printStackTrace();
            }
    }


    public void commitWrite()
    {
        if(iw!=null)
            try {
                iw.commit();
            } catch (IOException e) {
                e.printStackTrace();
            }
    }

    public void closeWrite()
    {
        if(iw!=null)
            try {
                iw.close();
                iw=null;
            } catch (IOException e) {
                e.printStackTrace();
            }
    }


    public void addImageDoc(BufferedImage img, String imageFilePath) {
        try {
            imageFilePath.replace("/work/scrapy/data/fruitimgs/1/","");
            Document document = globalDocumentBuilder.createDocument(img, imageFilePath);
            iw.addDocument(document);
        } catch (IOException e) {
            e.printStackTrace();
        }

    }




    public void addImageDoc(String imageFilePath) {
        System.out.println("path:"+imageFilePath);
        if(imageFilePath.contains(".DS_Store"))
            return;
        BufferedImage img = null;
        try {

            img = ImageIO.read(new FileInputStream(imageFilePath));

            addImageDoc(img, imageFilePath);
        }
        catch (Exception e) {
            e.printStackTrace();
            return;
        }
    }

    public List<Pic> searchSimlarImages(BufferedImage img) {
        List<Pic> imgs = new ArrayList<Pic>();
        ImageSearchHits hits = null;
        if(ir==null) initIndexReader(indexpath);
        try {
            hits = searcher.search(img, ir);
            for (int i = 0; i < hits.length(); i++) {
                String fileName = null;
                fileName = ir.document(hits.documentID(i)).getValues(DocumentBuilder.FIELD_NAME_IDENTIFIER)[0];
                // System.out.println(hits.score(i) + ": \t" + fileName);
                Pic pic=new Pic();
                pic.setPath(fileName);
                imgs.add(pic);
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
        return imgs;
    }


    public static void main(String args[]) {
        String path = "/work/scrapy/data/fruitimgs/1/fedab95498d86d1e890f96dd423eaeef9a769282.jpg";
        HDLireOper oper = new HDLireOper();
        oper.addImageDoc(path);
    }


}
